from anjuke_spider_enhanced import EnhancedAnjukeRentSpider
from config_enhanced import SUPPORTED_CITIES
import pandas as pd
import json
import time


class BatchProcessor:
    def __init__(self):
        self.spider = EnhancedAnjukeRentSpider()

    def load_tasks_from_json(self, config_file):
        """从JSON文件加载批量任务"""
        try:
            with open(config_file, 'r', encoding='utf-8') as f:
                tasks = json.load(f)
            return tasks
        except Exception as e:
            print(f"加载配置文件失败: {e}")
            return []

    def load_tasks_from_csv(self, csv_file):
        """从CSV文件加载批量任务"""
        try:
            df = pd.read_csv(csv_file)
            tasks = df.to_dict('records')
            return tasks
        except Exception as e:
            print(f"加载CSV文件失败: {e}")
            return []

    def process_batch(self, tasks, output_dir="batch_results"):
        """处理批量任务"""
        import os
        os.makedirs(output_dir, exist_ok=True)

        all_results = []

        for i, task in enumerate(tasks, 1):
            print(f"\n{'=' * 50}")
            print(f"处理任务 {i}/{len(tasks)}: {task.get('city', '未知')}-{task.get('district', '全市')}")
            print(f"{'=' * 50}")

            try:
                houses = self.spider.crawl(
                    city=task.get('city'),
                    district=task.get('district'),
                    min_price=task.get('min_price'),
                    max_price=task.get('max_price'),
                    room_type=task.get('room_type'),
                    source_type=task.get('source_type'),
                    max_pages=task.get('max_pages', 5),
                    get_details=task.get('get_details', False)
                )

                if houses:
                    # 保存单个任务结果
                    filename = f"{task['city']}_{task.get('district', '全市')}_{int(time.time())}.xlsx"
                    filepath = os.path.join(output_dir, filename)
                    self.spider.save_to_excel(houses, filepath)

                    # 添加到总结果
                    for house in houses:
                        house['任务来源'] = f"{task['city']}-{task.get('district', '全市')}"
                    all_results.extend(houses)

                # 任务间延迟
                if i < len(tasks):
                    delay = 10
                    print(f"等待 {delay} 秒后处理下一个任务...")
                    time.sleep(delay)

            except Exception as e:
                print(f"处理任务 {i} 时出错: {e}")
                continue

        # 保存合并结果
        if all_results:
            merged_file = os.path.join(output_dir, f"批量爬取合并结果_{int(time.time())}.xlsx")
            self.spider.save_to_excel(all_results, merged_file)
            print(f"\n批量处理完成！合并结果已保存到: {merged_file}")

        return all_results


def create_sample_config():
    """创建示例配置文件"""
    sample_tasks = [
        {
            "city": "广州",
            "district": "天河",
            "min_price": "2000",
            "max_price": "5000",
            "room_type": "二室",
            "max_pages": 3,
            "get_details": False
        },
        {
            "city": "广州",
            "district": "白云",
            "min_price": "1000",
            "max_price": "3000",
            "room_type": "一室",
            "max_pages": 3,
            "get_details": False
        },
        {
            "city": "深圳",
            "district": "南山",
            "min_price": "3000",
            "max_price": "8000",
            "room_type": "二室",
            "max_pages": 2,
            "get_details": True
        }
    ]

    with open('batch_config.json', 'w', encoding='utf-8') as f:
        json.dump(sample_tasks, f, ensure_ascii=False, indent=2)

    print("示例配置文件 'batch_config.json' 已创建")


if __name__ == "__main__":
    processor = BatchProcessor()

    print("批量处理工具")
    print("1. 创建示例配置文件")
    print("2. 执行批量爬取")

    choice = input("请选择: ").strip()

    if choice == "1":
        create_sample_config()
    elif choice == "2":
        config_file = input("配置文件路径 (默认: batch_config.json): ").strip() or "batch_config.json"
        tasks = processor.load_tasks_from_json(config_file)
        if tasks:
            processor.process_batch(tasks)
        else:
            print("没有找到有效的任务配置")